import numpy as np
from PIL import Image, ImageDraw
import os
import random

# 设置随机种子（可复现）
np.random.seed(42)
random.seed(42)


DESKTOP_PATH = os.path.join(os.path.expanduser("~"), "Desktop")
output_dir = os.path.join(DESKTOP_PATH, "CNN_model","glitch_patterns")

os.makedirs(output_dir, exist_ok=True)

WIDTH, HEIGHT = 1080, 2340 # 图像大小（根据手机屏幕分辨率）

def save_rgba_image(img_array, filename):
    """保存 RGBA 图像"""
    img = Image.fromarray(img_array, 'RGBA')
    img.save(os.path.join(output_dir, filename), "PNG")

def generate_gaussian_blur_patch(num_samples=3):
    """高斯模糊"""
    from scipy.ndimage import gaussian_filter

    for i in range(num_samples):
        # 透明背景 (RGBA)
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        # 随机位置、大小
        w, h = random.randint(100, 400), random.randint(100, 400)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        # 随机颜色（RGB），alpha=255
        color = (random.randint(100, 255), random.randint(0, 100), random.randint(0, 100))
        
        # 在区域内填充颜色
        img[y:y+h, x:x+w, 0] = color[0]
        img[y:y+h, x:x+w, 1] = color[1]
        img[y:y+h, x:x+w, 2] = color[2]
        img[y:y+h, x:x+w, 3] = 255  # 完全不透明

        # 提取 RGB 通道进行高斯模糊
        blurred_rgb = np.zeros((h, w, 3))
        for c in range(3):
            blurred_rgb[:, :, c] = gaussian_filter(img[y:y+h, x:x+w, c], sigma=15)
        
        # 替换回原图像
        img[y:y+h, x:x+w, :3] = blurred_rgb.astype(np.uint8)
        
        save_rgba_image(img, f"gaussian_blur_{i}.png")

def generate_color_blobs(num_samples=3):
    """随机小色块"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        draw = ImageDraw.Draw(Image.fromarray(img, 'RGBA'), 'RGBA')
        
        num_blobs = random.randint(5, 15)
        for _ in range(num_blobs):
            size = random.randint(10, 60)
            x = random.randint(0, WIDTH - size)
            y = random.randint(0, HEIGHT - size)
            color = (
                random.randint(150, 255),
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(100, 255)  # 随机透明度
            )
            draw.ellipse([x, y, x + size, y + size], fill=color)
        
        # 转换回 numpy array
        img = np.array(draw._image)
        save_rgba_image(img, f"color_blobs_{i}.png")

def generate_mosaic(num_samples=3):
    """马赛克"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        # 随机区域
        w, h = random.randint(120, 300), random.randint(120, 300)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        # 大小
        block_size = random.randint(8, 20)
        
        for by in range(y, y + h, block_size):
            for bx in range(x, x + w, block_size):
                # 随机颜色
                color = (
                    random.randint(0, 255),
                    random.randint(0, 255),
                    random.randint(0, 255),
                    255
                )
                # 填充块
                y_end = min(by + block_size, y + h)
                x_end = min(bx + block_size, x + w)
                img[by:y_end, bx:x_end, :] = color
        
        save_rgba_image(img, f"mosaic_{i}.png")

def generate_stripes(num_samples=3):
    """彩色条纹干扰"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        x = random.randint(0, WIDTH // 2)
        w = random.randint(100, 400)
        direction = random.choice(['vertical', 'horizontal'])
        
        if direction == 'vertical':
            spacing = random.randint(5, 20)
            for sx in range(x, x + w, spacing):
                color = (
                    random.randint(100, 255),
                    random.randint(0, 100),
                    random.randint(150, 255),
                    random.randint(100, 200)
                )
                img[:, sx:sx+spacing, :] = color
        else:
            y = random.randint(0, HEIGHT // 2)
            h = random.randint(100, 400)
            spacing = random.randint(5, 20)
            for sy in range(y, y + h, spacing):
                color = (
                    random.randint(150, 255),
                    random.randint(100, 200),
                    random.randint(0, 50),
                    random.randint(100, 200)
                )
                img[sy:sy+spacing, :, :] = color
        
        save_rgba_image(img, f"stripes_{i}.png")

def generate_salt_pepper_noise(num_samples=3):
    """椒盐噪点/黑白马赛克"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        w, h = random.randint(150, 400), random.randint(150, 400)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        # 生成随机噪点（黑白点）
        noise = np.random.choice([0, 128, 255], size=(h, w), p=[0.8, 0.1, 0.1])
        for c in range(3):  # RGB
            img[y:y+h, x:x+w, c] = noise
        img[y:y+h, x:x+w, 3] = 255  # alpha 完全不透明
        
        save_rgba_image(img, f"salt_pepper_{i}.png")

def generate_color_tint(num_samples=3):
    """色偏区域"""
    tints = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
    
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        w, h = random.randint(200, 500), random.randint(200, 500)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        tint = random.choice(tints)
        alpha = random.randint(50, 150)  # 半透明
        
        for c in range(3):
            img[y:y+h, x:x+w, c] = tint[c]
        img[y:y+h, x:x+w, 3] = alpha  # 半透明
        
        save_rgba_image(img, f"color_tint_{i}.png")

def generate_lines_interference(num_samples=3):
    """竖线或横线干扰（彩色或黑白）"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        direction = random.choice(['vertical', 'horizontal'])
        line_count = random.randint(10, 50)
        thickness = random.randint(1, 3)
        intensity = random.randint(150, 255)
        alpha = random.randint(100, 255)

        if direction == 'vertical':
            spacing = WIDTH // line_count
            for x in range(0, WIDTH, spacing):
                color = (
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    alpha
                )
                x_end = x + thickness
                if x_end > WIDTH: x_end = WIDTH
                img[:, x:x_end, :] = color
        else:  # horizontal
            spacing = HEIGHT // line_count
            for y in range(0, HEIGHT, spacing):
                color = (
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    random.choice([0, intensity, 255]) if random.random() > 0.5 else intensity,
                    alpha
                )
                y_end = y + thickness
                if y_end > HEIGHT: y_end = HEIGHT
                img[y:y_end, :, :] = color

        save_rgba_image(img, f"lines_interference_{i}.png")


def generate_screen_tearing(num_samples=3):
    """屏幕撕裂"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        # 位置
        tear_y = random.randint(200, HEIGHT - 200)
        offset = random.randint(20, 100) * random.choice([-1, 1])  # 左右偏移

        # 区域高度
        tear_height = random.randint(30, 100)
        
        # 随机颜色块
        for dy in range(tear_y, tear_y + tear_height):
            color = (
                random.randint(100, 255),
                random.randint(50, 150),
                random.randint(50, 150),
                random.randint(100, 200)
            )
            # 偏移后的起始列
            start_x = max(0, -offset)
            end_x = min(WIDTH, WIDTH - offset) if offset > 0 else WIDTH
            img[dy, start_x:end_x, :] = color

        save_rgba_image(img, f"screen_tearing_{i}.png")


def generate_color_inversion(num_samples=3):
    """局部反色（负片效果）"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        w, h = random.randint(150, 400), random.randint(150, 400)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        # 反色区域：RGB 取反，Alpha 全不透明
        for c in range(3):
            img[y:y+h, x:x+w, c] = 255  # 取反后的值（叠加时与底图 XOR）
        img[y:y+h, x:x+w, 3] = random.randint(80, 150)  # 半透明，表示“反色层”

        # 注意：反色效果需在叠加时实现（见下方说明）
        save_rgba_image(img, f"color_inverted_region_{i}.png")


def generate_flickering_block(num_samples=3):
    """闪烁区域"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        w, h = random.randint(100, 300), random.randint(100, 300)
        x = random.randint(0, WIDTH - w)
        y = random.randint(0, HEIGHT - h)
        
        # 使用亮色闪烁（白/黄）
        flicker_color = random.choice([
            (255, 255, 0, 200),  # 黄
            (255, 200, 0, 180),
            (255, 255, 255, 150)
        ])
        
        # 用随机斑点模拟不稳定闪烁
        for _ in range(20):
            rx = random.randint(x, x + w - 1)
            ry = random.randint(y, y + h - 1)
            size = random.randint(5, 20)
            for dy in range(ry, min(ry + size, HEIGHT)):
                for dx in range(rx, min(rx + size, WIDTH)):
                    img[dy, dx] = flicker_color

        save_rgba_image(img, f"flickering_block_{i}.png")


def generate_dead_pixels(num_samples=3):
    """坏点阵列（黑色或固定颜色像素点）"""
    for i in range(num_samples):
        img = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
        
        num_dots = random.randint(20, 100)
        
        for _ in range(num_dots):
            x = random.randint(0, WIDTH - 1)
            y = random.randint(0, HEIGHT - 1)
            
            stuck_color = random.choice([
                (0, 0, 0, 255),        # 黑点
                (255, 0, 0, 255),      # 红点
                (0, 255, 0, 255),      # 绿点
                (0, 0, 255, 255),      # 蓝点
            ])
            img[y, x] = stuck_color

        save_rgba_image(img, f"dead_pixels_{i}.png")


def generate_all_patterns(num_per_type=3):
    print("Generating Gaussian Blur Patches...")
    generate_gaussian_blur_patch(num_per_type)
    
    #print("Generating Color Blobs...")
    #generate_color_blobs(num_per_type)
    
    print("Generating Mosaic Patterns...")
    generate_mosaic(num_per_type)
    
    print("Generating Stripes...")
    generate_stripes(num_per_type)
    
    print("Generating Salt-and-Pepper Noise...")
    generate_salt_pepper_noise(num_per_type)
    
    print("Generating Color Tint Patterns...")
    generate_color_tint(num_per_type)
    

    print("Generating Lines Interference...")
    generate_lines_interference(num_per_type)
    
    
    print("Generating Screen Tearing...")
    generate_screen_tearing(num_per_type)
    
    #print("Generating Color Inversion Regions...")
    #generate_color_inversion(num_per_type)
    
    #print("Generating Flickering Blocks...")
    #generate_flickering_block(num_per_type)
    
    #print("Generating Dead Pixels...")
    #generate_dead_pixels(num_per_type)
    
    print(f"All {num_per_type} samples per type generated and saved to '{output_dir}/'")


# 运行生成
if __name__ == "__main__":
    generate_all_patterns(num_per_type=5)  # 每类生成5个示例
